# -*- coding: utf-8 -*-
"""
Solve Sudoku Problem

@author: LiuHuan
"""

import numpy as np
from typing import List, Tuple, Optional

set1_9 = set(range(1, 10))


class Node():
    """Sudoku Node"""

    def __init__(self, problem: np.ndarray):
        if isinstance(problem, list) and isinstance(problem[0], list):
            self.problem = np.array(problem, dtype=np.int)
        elif isinstance(problem, np.ndarray) and len(problem.shape) == 2:
            self.problem = problem
        else:
            self.problem = problem
            raise ValueError("Unsupported Problem Type")
        self.parent = None
        self.fill_num = 0
        self.fill_loc = (0, 0)
        self.solutions = []

    def opration(self) -> str:
        return "fill in {} at {},{}".format(
            self.fill_num, self.fill_loc[0] + 1, self.fill_loc[1] + 1
        )


class Sudoku(Node):
    def expand(
        self
    ) -> Tuple[
        str,
        Optional[List[List[Tuple[int, int]]]],
        Optional[List[List[int]]]
    ]:
        min_num = np.inf
        sel_index = (0, 0)
        sel_candidate = []
        blank_index = np.argwhere(self.problem == 0)
        if len(blank_index) == 0:
            return "Solution", None, None  # End
        for row, col in blank_index:
            row_set = set(self.problem[row, :])
            col_set = set(self.problem[:, col])
            row_area = row // 3
            col_area = col // 3
            area_set = set(
                self.problem[
                    row_area*3:row_area*3+3, col_area*3:col_area*3+3
                ].ravel()
            )

            candidate = sorted(
                (set1_9 - row_set) & (set1_9 - col_set) & (set1_9 - area_set),
                reverse=True
            )

            if candidate and len(candidate) < min_num:
                min_num = len(candidate)
                sel_index = (row, col)
                sel_candidate = candidate
        if np.isinf(min_num):
            return "NoSolution", None, None
        return "Candidate", [sel_index] * len(sel_candidate), sel_candidate

    def solve(self):
        print("Searching...", end=" ")
        close_table = set()
        close_table.add(tuple(self.problem.ravel()))
        seq = [self]
        solutions = []
        while True:
            if not seq:  # End
                print("End. Expand %d nodes" % (len(close_table)))
                break
            node = seq.pop()
            state, sel_index_list, candidate = node.expand()
            close_table.add(tuple(node.problem.ravel()))
            if state == "Solution":  # found one solution
                solution = []
                tmp_node = node
                while True:
                    solution.append(tmp_node)
                    tmp_node = tmp_node.parent
                    if tmp_node is None:
                        break
                solution = solution[::-1]
                solutions.append(solution)
            elif state == "NoSolution":
                continue
            else:
                for i in range(len(candidate)):
                    tmp = node.problem.copy()
                    tmp[sel_index_list[i]] = candidate[i]
                    if tuple(tmp.ravel()) not in close_table:
                        child_node = type(node)(tmp)
                        child_node.parent = node
                        child_node.fill_loc = sel_index_list[i]
                        child_node.fill_num = candidate[i]
                        seq.append(child_node)
        self.solutions = solutions
        return self

    def print_final_results(self):
        print("====== All final results: ======")
        for i in range(len(self.solutions)):
            print(i+1, ":")
            print(self.solutions[i][-1].problem)
        print("======        End.        ======")
        return self

    def print_detailed_steps(self):
        print("======   Detailed Steps   ======")
        for i, solution in enumerate(self.solutions):
            print("----- {}-th solution begin -----".format(i+1))
            print("Solution:")
            print(solution[0].problem)
            for node_i in solution[1:]:
                print(node_i.opration())
                print(node_i.problem)
            print("----- {}-th solution  end  -----".format(i+1))
        print("======        End.        ======")
        return self


class FastSudoku(Sudoku):
    def expand(
        self
    ) -> Tuple[
        str,
        Optional[List[List[Tuple[int, int]]]],
        Optional[List[List[int]]]
    ]:
        count = [0] * 10
        for i in self.problem.ravel():
            count[i] += 1
        if count[0] == 0:
            return "Solution", None, None  # End
        find_Flag = False
        d = {}
        sel_index_list_list = []
        candidate_list = []
        for i in range(1, 10):
            state, sel_index_list, candidate = self._fill_i(i)
            sel_index_list_list.append(sel_index_list)
            candidate_list.append(candidate)
            if state == "Candidate":
                d[i] = len(candidate)
                find_Flag = True
        if not find_Flag:
            return "NoSolution", None, None
        sel_i, _ = sorted(d.items(), key=lambda x: x[1])[0]
        return "Candidate", sel_index_list_list[sel_i-1], candidate_list[sel_i-1]

    def _fill_i(
        self,
        sel_i: int
    ) -> Tuple[
        str,
        List[Tuple[int, int]],
        List[int]
    ]:
        sel_index_list = []
        candidate = []
        not_i = np.ones((9, 9), dtype=np.bool)
        for row, col in np.argwhere(self.problem == sel_i):
            not_i[row, :] = 0
            not_i[:, col] = 0
            row_area = row // 3
            col_area = col // 3
            not_i[row_area*3:row_area*3+3, col_area*3:col_area*3+3] = 0
        not_i = not_i & np.logical_not(self.problem.astype(np.bool))
        area_list = []
        for row_area in range(3):
            for col_area in range(3):
                area_count = np.sum(not_i[row_area*3:row_area*3+3, col_area*3:col_area*3+3])
                if area_count > 0:
                    area_list.append((
                        (row_area, col_area),
                        area_count
                    ))
        if area_list:
            area_sorted = sorted(area_list, key=lambda x: x[1])
            row_area, col_area = area_sorted[0][0]
            area_count = area_sorted[0][1]
            for row_i, col_i in np.argwhere(not_i[row_area*3:row_area*3+3, col_area*3:col_area*3+3]):
                sel_row = row_i + row_area * 3
                sel_col = col_i + col_area * 3
                sel_index_list.append((sel_row, sel_col))
                candidate.append(sel_i)
            return "Candidate", sel_index_list, candidate
        if not sel_index_list:
            return "NoSolution", None, None


if __name__ == "__main__":

    # problem = [
    #     [0, 0, 7, 0, 0, 0, 0, 0, 8],
    #     [0, 4, 0, 7, 0, 0, 1, 2, 0],
    #     [2, 3, 0, 0, 0, 1, 6, 0, 0],
    #     [9, 0, 0, 0, 1, 0, 0, 0, 0],
    #     [0, 0, 4, 9, 7, 8, 3, 0, 0],
    #     [0, 0, 0, 0, 4, 0, 0, 0, 5],
    #     [0, 0, 9, 8, 0, 0, 0, 5, 1],
    #     [0, 1, 3, 0, 0, 9, 0, 7, 0],
    #     [7, 0, 0, 0, 0, 0, 8, 0, 0]
    # ]

    # problem = [
    #     [5, 0, 4, 1, 0, 0, 6, 0, 0],
    #     [0, 0, 0, 0, 0, 4, 0, 3, 7],
    #     [0, 2, 3, 0, 0, 0, 0, 0, 9],
    #     [0, 8, 2, 4, 0, 0, 0, 0, 1],
    #     [9, 0, 0, 0, 7, 0, 0, 0, 3],
    #     [7, 0, 0, 0, 0, 1, 8, 9, 0],
    #     [2, 0, 0, 0, 0, 0, 3, 8, 0],
    #     [6, 1, 0, 3, 0, 0, 0, 0, 0],
    #     [0, 0, 5, 0, 0, 7, 2, 0, 6]
    # ]

    problem = [
        [8, 7, 0, 0, 0, 0, 0, 0, 3],
        [0, 9, 4, 0, 0, 0, 0, 6, 5],
        [0, 0, 2, 0, 1, 0, 0, 0, 0],
        [0, 6, 1, 0, 0, 3, 0, 0, 0],
        [0, 4, 0, 1, 0, 0, 0, 7, 0],
        [0, 0, 0, 4, 0, 0, 3, 2, 0],
        [0, 0, 0, 0, 9, 0, 6, 0, 0],
        [7, 2, 0, 0, 0, 0, 8, 5, 0],
        [4, 0, 0, 0, 0, 0, 0, 3, 9]
    ]

    # solver1 = Sudoku(problem).solve()
    solver2 = FastSudoku(problem).solve()
    # solver1.print_final_results()
    solver2.print_final_results()
    # solver1.print_detailed_steps()
    solver2.print_detailed_steps()
